import os
import argparse
import numpy as np
from typing import List, Dict, Any
from sentence_transformers import SentenceTransformer
import faiss
import json

class LocalTextEmbedder:
    """使用本地大模型生成文本向量的类"""
    
    def __init__(self, model_name: str = "all-MiniLM-L6-v2", device: str = "cpu"):
        """
        初始化文本向量化器
        
        Args:
            model_name: 本地模型名称或路径
            device: 运行模型的设备，如"cpu"或"cuda"
        """
        self.model = SentenceTransformer(model_name, device=device)
        self.index = None
        self.texts = []
        self.embeddings = []
    
    def embed_texts(self, texts: List[str]) -> np.ndarray:
        """
        将文本列表转换为向量表示
        
        Args:
            texts: 文本列表
            
        Returns:
            向量表示的numpy数组
        """
        return self.model.encode(texts, convert_to_numpy=True)
    
    def build_index(self, texts: List[str], save_path: str = None) -> None:
        """
        构建向量索引并保存文本和向量
        
        Args:
            texts: 文本列表
            save_path: 保存索引的路径
        """
        self.texts = texts
        self.embeddings = self.embed_texts(texts)
        dimension = self.embeddings.shape[1]
        
        # 使用FAISS构建索引
        self.index = faiss.IndexFlatL2(dimension)
        self.index.add(self.embeddings)
        
        if save_path:
            self.save_index(save_path)
    
    def save_index(self, path: str) -> None:
        """保存索引和相关数据"""
        if self.index is None:
            raise ValueError("索引未构建，请先调用build_index方法")
        
        # 保存FAISS索引
        faiss.write_index(self.index, f"{path}.index")
        
        # 保存文本和嵌入向量
        data = {
            "texts": self.texts,
            "embeddings": self.embeddings.tolist()
        }
        
        with open(f"{path}.json", "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
    
    def load_index(self, path: str) -> None:
        """加载已保存的索引和相关数据"""
        # 加载FAISS索引
        self.index = faiss.read_index(f"{path}.index")
        
        # 加载文本和嵌入向量
        with open(f"{path}.json", "r", encoding="utf-8") as f:
            data = json.load(f)
        
        self.texts = data["texts"]
        self.embeddings = np.array(data["embeddings"])
    
    def search_similar(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
        """
        搜索最相似的文本
        
        Args:
            query: 查询文本
            top_k: 返回的相似文本数量
            
        Returns:
            包含相似文本及其相似度的字典列表
        """
        if self.index is None:
            raise ValueError("索引未构建，请先调用build_index方法")
        
        query_embedding = self.embed_texts([query])
        distances, indices = self.index.search(query_embedding, top_k)
        
        results = []
        for i, idx in enumerate(indices[0]):
            if idx != -1:  # 排除无效索引
                results.append({
                    "text": self.texts[idx],
                    "similarity": 1 - distances[0][i] / 2,  # 将L2距离转换为相似度
                    "index": idx
                })
        
        return results


def main():
    parser = argparse.ArgumentParser(description="本地文本向量化工具")
    parser.add_argument("--model", type=str, default="all-MiniLM-L6-v2", 
                        help="使用的模型名称或路径")
    parser.add_argument("--device", type=str, default="cpu", 
                        help="运行模型的设备，如'cpu'或'cuda'")
    parser.add_argument("--input_file", type=str, 
                        help="输入文本文件路径，每行一个文本")
    parser.add_argument("--output_path", type=str, 
                        help="保存索引的路径")
    parser.add_argument("--load_path", type=str, 
                        help="加载已有索引的路径")
    parser.add_argument("--query", type=str, 
                        help="查询相似文本")
    parser.add_argument("--top_k", type=int, default=5, 
                        help="返回的相似文本数量")
    
    args = parser.parse_args()
    
    # 初始化向量化器
    embedder = LocalTextEmbedder(model_name=args.model, device=args.device)
    print(f"使用模型: {args.model}")
    
    if args.load_path:
        # 加载已有索引
        embedder.load_index(args.load_path)
        print(f"已加载索引，包含 {len(embedder.texts)} 个文本")
    
    if args.input_file:
        # 从文件读取文本
        with open(args.input_file, "r", encoding="utf-8") as f:
            texts = [line.strip() for line in f if line.strip()]
        
        # 构建索引
        embedder.build_index(texts, args.output_path)
        print(f"已构建索引，包含 {len(texts)} 个文本")
    
    if args.query:
        # 搜索相似文本
        results = embedder.search_similar(args.query, args.top_k)
        
        print(f"\n查询: {args.query}")
        print("最相似的文本:")
        for i, result in enumerate(results):
            print(f"{i+1}. 相似度: {result['similarity']:.4f}")
            print(f"   文本: {result['text']}")
            print()


if __name__ == "__main__":
    main()    